package io.elegant.sessions

import android.content.Context
import android.media.AudioFormat
import android.media.AudioTrack
import android.media.MediaCodec
import android.media.MediaFormat
import com.google.protobuf.Any
import io.elegant.*
import io.elegant.ElegantApiGrpc.ElegantApiStub
import io.grpc.stub.StreamObserver
import kotlinx.coroutines.*
import java.nio.ByteBuffer


class AudioTrackSession(
    stub: ElegantApiStub,
    context: Context
): BaseSession(stub, context) {
    override val sendStreamObserver: StreamObserver<Elegant.Interaction> = stub.audioTrack(receiveStreamObserver)
    private var track: AudioTrack? = null
    private var mediaCodec: MediaCodec? = null
    override fun onCreate(sessionId: Long, futureId: Long, vararg argv: Any) {
        val sampleRate = if (argv.isNotEmpty()) argv[0].toIntValue() else 16000
        val channels = if (argv.size > 1) argv[1].toIntValue() else 1
        val channelConfig = if (channels == 2) AudioFormat.CHANNEL_OUT_STEREO else AudioFormat.CHANNEL_OUT_MONO
        val bufferSize = AudioTrack.getMinBufferSize(sampleRate, channelConfig, AudioFormat.ENCODING_PCM_16BIT)
        val format = AudioFormat.Builder()
            .setEncoding(AudioFormat.ENCODING_PCM_16BIT)
            .setChannelMask(channelConfig)
            .build()
        track = AudioTrack.Builder()
            .setAudioFormat(format)
            .setBufferSizeInBytes(bufferSize)
            .build()
        val mediaFormat = MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_OPUS, sampleRate, channels)
        val csd0bytes = byteArrayOf(
            // Opus
            0x4f, 0x70, 0x75, 0x73,
            // Head
            0x48, 0x65, 0x61, 0x64,
            // Version
            0x01,
            // Channel Count
            channels.toByte(),
            // Pre skip
            0x00, 0x00, 0x80.toByte(), 0xbb.toByte(), 0x00, 0x00,
            // Output Gain (Q7.8 in dB)
            0x00, 0x00,
            // Mapping Family
            0x00
        )
        val csd1bytes = byteArrayOf(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00)
        val csd2bytes = byteArrayOf(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00)
        val csd0 = ByteBuffer.wrap(csd0bytes)
        mediaFormat.setByteBuffer("csd-0", csd0)
        val csd1 = ByteBuffer.wrap(csd1bytes)
        mediaFormat.setByteBuffer("csd-1", csd1)
        val csd2 = ByteBuffer.wrap(csd2bytes)
        mediaFormat.setByteBuffer("csd-2", csd2)
        mediaCodec = MediaCodec.createDecoderByType(MediaFormat.MIMETYPE_AUDIO_OPUS).apply {
            configure(mediaFormat, null, null, 0)
            start()
        }
        super.onCreate(sessionId, futureId, *argv)
    }
    override fun onDestroy() {
        mediaCodec?.release()
        mediaCodec = null
        track?.release()
        track = null
        super.onDestroy()
    }
    override suspend fun onReceive(name: String, sessionId: Long, futureId: Long, vararg argv: Any) {
        when (name) {
            "write" -> write(sessionId, futureId, *argv)
            "stop" -> stop()
            "play" -> play()
            "pause" -> pause()
            "setSpeed" -> setSpeed(*argv)
            "setVolume" -> setVolume(*argv)
            else -> super.onReceive(name, sessionId, futureId, *argv)
        }
    }
    private suspend fun write(sessionId: Long, futureId: Long, vararg argv: Any) {
        if (argv.isEmpty())
            throw IllegalArgumentException("Missing param `data` from the server calling this `write` method.")
        val data = argv[0].toBytesValue()
        val pcm = csIo.async {
            val timeoutUs = 1000L
            try {
                mediaCodec
                    ?.dequeueInputBuffer(timeoutUs)
                    ?.takeIf { it >= 0 }
                    ?.let {
                        mediaCodec?.getInputBuffer(it)?.let { it1 ->
                            it1.clear()
                            it1.put(data, 0, data.size)
                        }
                        mediaCodec?.queueInputBuffer(it, 0, data.size, 0, 0)
                    }
            } catch (e: IllegalStateException) {
                e.printStackTrace()
            }
            val decodeBufferInfo = MediaCodec.BufferInfo()
            try {
                mediaCodec
                    ?.dequeueOutputBuffer(decodeBufferInfo, timeoutUs)
                    ?.takeIf { it >= 0 }
                    ?.let {
                        val out = ByteArray(decodeBufferInfo.size)
                        mediaCodec?.getOutputBuffer(it)?.let { it1 ->
                            it1.get(out)
                            it1.clear()
                        }
                        mediaCodec?.releaseOutputBuffer(it, false)
                        out
                    }
            } catch (e: IllegalStateException) {
                e.printStackTrace()
                null
            }
        }.await() ?: byteArrayOf()
        track?.let {
            onWrite(true.toAny(), sessionId, futureId)
            it.write(pcm, 0, pcm.size)
        } ?: onWrite(false.toAny(), sessionId, futureId)
    }
    private fun play() = track?.play()
    private fun stop() = track?.stop()
    private fun pause() = track?.pause()
    private fun setSpeed(vararg argv: Any) = track?.let {
        if (argv.isEmpty())
            throw IllegalArgumentException("Missing param `rate` from the server calling this `setSpeed` method.")
        it.playbackParams = it.playbackParams.also { it1 ->
            it1.speed = argv[0].toFloatValue()
        }
    }
    private fun setVolume(vararg argv: Any) = track?.let {
        if (argv.isEmpty())
            throw IllegalArgumentException("Missing param `volume` from the server calling this `setVolume` method.")
        val left = argv[0].toFloatValue()
        val right = if (argv.size > 1) argv[1].toFloatValue() else left
        @Suppress("DEPRECATION")
        it.setStereoVolume(left, right)
    }
    private fun onWrite(result: Any, sessionId: Long, futureId: Long) = sendAction("onWrite", sessionId, futureId, result)
}